package com.demo.config;
import com.demo.util.ShardingAlgorithmUtil;
import com.google.common.collect.Range;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang.time.DateUtils;
import org.apache.shardingsphere.sharding.api.sharding.standard.PreciseShardingValue;
import org.apache.shardingsphere.sharding.api.sharding.standard.RangeShardingValue;
import org.apache.shardingsphere.sharding.api.sharding.standard.StandardShardingAlgorithm;

import java.text.DateFormat;
import java.text.SimpleDateFormat;
import java.util.*;
import java.util.stream.Collectors;

/**
 * 订单时间分片算法
 */


@Slf4j
public class OrderTimeShardingAlgorithm implements StandardShardingAlgorithm<Date> {
    private static final DateFormat TABLE_SHARD_TIME_FORMAT = new SimpleDateFormat("yyyyMM");

    // 表分片符号。如：tb_order_202407
    private static final String TABLE_SPLIT_SYMBOL = "_";

    private Properties props;

    @Override
    public String doSharding(Collection<String> availableTargetNames, PreciseShardingValue<Date> shardingValue) {
        String logicTableName = shardingValue.getLogicTableName();
        log.info("精准分片，逻辑表名：{}，节点表名：{}", logicTableName, availableTargetNames);
        Date time = shardingValue.getValue();
        String result = logicTableName + TABLE_SPLIT_SYMBOL + TABLE_SHARD_TIME_FORMAT.format(time);
        // 在配置中，只配置了逻辑表名。如果只有一个，且是逻辑表名，说明需要获取所有表名
        initAvailableTargetNames(availableTargetNames, logicTableName);
        return getAndCreateShardingTable(logicTableName, result, availableTargetNames);
    }

    @Override
    public Collection<String> doSharding(Collection<String> availableTargetNames, RangeShardingValue<Date> shardingValue) {
        String logicTableName = shardingValue.getLogicTableName();
        log.info("精准分片，逻辑表名：{}，节点表名：{}", logicTableName, availableTargetNames);

        // 在配置中，只配置了逻辑表名。如果只有一个，且是逻辑表名，说明需要获取所有表名
        initAvailableTargetNames(availableTargetNames, logicTableName);

        Range<Date> valueRange = shardingValue.getValueRange();

        // 如果没有最大值或最小值，则全库扫描
        if(!valueRange.hasLowerBound() || !valueRange.hasUpperBound()) {
            return availableTargetNames;
        }

        Date min = valueRange.lowerEndpoint();
        Date max = valueRange.upperEndpoint();

        Set<String> rs = new HashSet<>();
        while (min.compareTo(max) <= 0) {
            String tableName = logicTableName + "_" + TABLE_SHARD_TIME_FORMAT.format(min);
            rs.add(tableName);
            min = DateUtils.addMonths(min, 1);
        }
        return getAndCreateShardingTable(logicTableName, rs, availableTargetNames);
    }

    private void initAvailableTargetNames(Collection<String> availableTargetNames, String logicTableName) {
        if(availableTargetNames.size() == 1 && availableTargetNames.contains(logicTableName)) {
            Set<String> allTableNameBySchema = ShardingAlgorithmUtil.getAllTableNameBySchema(logicTableName);
            availableTargetNames.clear();
            availableTargetNames.addAll(allTableNameBySchema);
        }
    }

    /**
     * 检查可用的真实表，如果表名不存在，则创建新表
     * @param logicTableName 逻辑表
     * @param resultTableNames 操作需要的真实表
     * @param availableTargetNames 可用的真实表
     * @return 分片的真实表
     */
    private List<String> getAndCreateShardingTable(String logicTableName, Set<String> resultTableNames, Collection<String> availableTargetNames) {
        return resultTableNames.stream().map(name -> getAndCreateShardingTable(logicTableName, name, availableTargetNames)).collect(Collectors.toList());
    }

    /**
     * 检查可用的真实表，如果表名不存在，则创建新表
     * @param logicTableName
     * @param resultTableName
     * @param availableTargetNames
     * @return
     */
    private String getAndCreateShardingTable(String logicTableName, String resultTableName, Collection<String> availableTargetNames) {
        if(availableTargetNames.contains(resultTableName)) {
            return resultTableName;
        }
        boolean rs = ShardingAlgorithmUtil.createShardingTable(logicTableName, resultTableName);
        if(rs) {
            availableTargetNames.add(resultTableName);
            return resultTableName;
        }
        return null;
    }

    @Override
    public Properties getProps() {
        return props;
    }

    @Override
    public void init(Properties properties) {
        this.props = properties;
    }

}
